## 负责展示所有集合，并能够删除指定的集合

# 列出所有Milvus集合
# 提供删除指定集合的功能
# 操作简单直观的菜单界面

from config import envConfig
from pymilvus import connections, utility
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('manage_collections')


def connect_milvus():
    """连接Milvus数据库"""
    try:
        connections.connect(
            alias="default",
            host=envConfig.MILVUS_HOST,
            port=envConfig.MILVUS_PORT,
            user=envConfig.MILVUS_USER,
            password=envConfig.MILVUS_PASSWORD,
            secure=False
        )
        return True
    except Exception as e:
        logger.error(f"Milvus连接失败: {str(e)}")
        return False


def list_collections():
    """列出所有集合"""
    if not connect_milvus():
        return []

    try:
        collections = utility.list_collections()
        if not collections:
            print("Milvus 中没有集合")
            return []

        print("\nMilvus 中的集合列表:")
        for i, col in enumerate(collections, 1):
            print(f"{i}. {col}")
        return collections
    except Exception as e:
        logger.error(f"获取集合列表失败: {str(e)}")
        return []


def delete_collection(collection_name):
    """删除指定集合"""
    if not connect_milvus():
        return False

    try:
        if utility.has_collection(collection_name):
            confirm = input(f"确定要删除集合 '{collection_name}' 吗? (y/n): ").strip().lower()
            if confirm == "y":
                utility.drop_collection(collection_name)
                print(f"集合 '{collection_name}' 已成功删除")
                return True
            else:
                print("删除操作已取消")
                return False
        else:
            print(f"集合 '{collection_name}' 不存在")
            return False
    except Exception as e:
        logger.error(f"删除集合失败: {str(e)}")
        return False


def main():
    """主函数：集合管理界面"""
    print("===== Milvus 集合管理器 =====")

    while True:
        print("\n选项:")
        print("1. 列出所有集合")
        print("2. 删除集合")
        print("3. 退出")

        choice = input("请选择操作: ").strip()

        if choice == "1":
            list_collections()
        elif choice == "2":
            collections = list_collections()
            if not collections:
                print("没有可删除的集合")
                continue

            try:
                index = int(input("输入要删除的集合编号: ").strip())
                if 1 <= index <= len(collections):
                    collection_name = collections[index - 1]
                    delete_collection(collection_name)
                else:
                    print("无效的编号，请重新输入")
            except ValueError:
                print("请输入有效的数字")
        elif choice == "3":
            print("退出集合管理器")
            break
        else:
            print("无效的选择，请重新输入")


if __name__ == "__main__":
    main()